# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.

import argparse
import fire

from llama import Llama

def main(args):
    ckpt_dir = args.ckpt_dir
    tokenizer_path = args.tokenizer_path
    temperature = args.temperature
    top_p = args.top_p
    max_seq_len = args.max_seq_len
    max_batch_size = args.max_batch_size
    max_gen_len = args.max_gen_len
    
    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
    )

    prompts = [
        # For these prompts, the expected answer is the natural continuation of the prompt
        "I believe the meaning of life is",
        "Simply put, the theory of relativity states that ",
        """A brief message congratulating the team on the launch:

        Hi everyone,
        
        I just """,
        # Few shot prompt (providing a few examples before asking model to complete more);
        """Translate English to French:
        
        sea otter => loutre de mer
        peppermint => menthe poivrée
        plush girafe => girafe peluche
        cheese =>""",
    ]
    results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )
    for prompt, result in zip(prompts, results):
        print(prompt)
        print(f"> {result['generation']}")
        print("\n==================================\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--ckpt_dir', default="", type=str, required=True)
    parser.add_argument('--tokenizer_path', default="", type=str, required=True)
    parser.add_argument('--temperature', default=0.6, type=float)
    parser.add_argument('--top_p', default=0.9, type=float)
    parser.add_argument('--max_seq_len', default=512, type=int)
    parser.add_argument('--max_batch_size', default=4, type=int)
    parser.add_argument('--max_gen_len', default=100, type=int)
    parser.add_argument('--local_rank', default=0, type=int)
    args = parser.parse_args()
    
    main(args)
